import backtrader as bt
import pandas as pd
import akshare as ak
from datetime import datetime

# 自定义策略
class MyStrategy(bt.Strategy):
    params = (
        ('fast_period', 5),
        ('slow_period', 20),
    )
    
    def __init__(self):
        self.data_close = self.datas[0].close
        # 添加快速和慢速移动平均线指标
        self.fast_sma = bt.indicators.SimpleMovingAverage(self.datas[0], period=self.params.fast_period)
        self.slow_sma = bt.indicators.SimpleMovingAverage(self.datas[0], period=self.params.slow_period)
        self.cross_up = bt.indicators.CrossUp(self.fast_sma, self.slow_sma)  # 金叉信号
        self.cross_down = bt.indicators.CrossDown(self.fast_sma, self.slow_sma)  # 死叉信号

    def next(self):
        if not self.position:  # 如果没有持仓
            if self.cross_up[0] > 0:  # 当快线由下向上穿过慢线，即金叉，买入
                self.buy()
        else:  # 如果已经持有仓位
            if self.cross_down[0] > 0:  # 当慢线由下向上穿过快线，即死叉，卖出
                self.close()

# 设置回测参数
start_date = datetime(2020, 1, 1)
end_date = datetime(2024, 1, 1)
start_cash = 100000
commission = 0.002
#===================================================================
# 股票代码列表
index_stock_cons_df = ak.index_stock_cons(symbol="000016")

# 创建一个空列表，用于存储股票代码
symbols = []
for stock in index_stock_cons_df['品种代码']:
    symbols.append(stock)

#======================================================================

def run_backtest_for_stocks(symbols):
    
    resultsEx = [] #excel
    
    for symbol in symbols:
        # 获取股票历史数据
        try:
            stock_hfq_df = ak.stock_zh_a_hist(symbol=symbol, adjust="hfq")
            if stock_hfq_df is None or stock_hfq_df.empty:
                print(f"No data found for {symbol}. Skipping...")
                continue
            
            stock_hfq_df = stock_hfq_df.iloc[:, :6]
            stock_hfq_df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
            stock_hfq_df['date'] = pd.to_datetime(stock_hfq_df['date'])
            stock_hfq_df.set_index('date', inplace=True)
            
            # 筛选指定日期范围内的数据
            data = stock_hfq_df[(stock_hfq_df.index >= start_date) & (stock_hfq_df.index <= end_date)]
        except Exception as e:
            print(f"Error fetching data for {symbol}: {e}")
            continue
        
        cerebro = bt.Cerebro()
        
        # 将数据添加到Cerebro
        data_feed = bt.feeds.PandasData(dataname=data)
        cerebro.adddata(data_feed)
        
        cerebro.addstrategy(MyStrategy, fast_period=5, slow_period=20)
        cerebro.broker.setcash(start_cash)
        cerebro.broker.setcommission(commission=commission)
        
        # 运行回测
        results = cerebro.run(tradehistory=True)
        
        # 打印结果
        if results:
            final_value = results[0].broker.getvalue()
            pnl = final_value - start_cash
            print(f"{symbol}: Final Portfolio Value: {final_value:.2f}, Net Profit: {pnl:.2f}")
            
            result_item = {'代码': symbol, 'Portfolio':final_value, 'Profit':pnl} 
            resultsEx.append(result_item) 
    
    # 将结果列表转换为pandas DataFrame  
    df = pd.DataFrame(resultsEx)  
  
    # 将DataFrame保存到Excel文件  
    excel_file = 'output.xlsx'  
    df.to_excel(excel_file, index=False, engine='openpyxl')  
  
    print(f"结果已保存到 {excel_file}")

# 执行回测
run_backtest_for_stocks(symbols)